package com.witmore.neutron.ai.cognitive.openai;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.witmore.neutron.ai.api.enums.ChatClientType;
import com.witmore.neutron.ai.db.entity.ModelConfigDO;
import com.witmore.neutron.ai.db.mapper.ModelConfigMapper;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Description: AiClientStart
 * Author: 豆子高
 * Date: 2025/6/3 21:22
 */
@Slf4j
@Component
public class AiClientModelCache {

    @Autowired
    private ModelConfigMapper modelConfigMapper;

    private static final Map<String, OpenAiChatModel> modelConfigMap = new HashMap<>();

    public OpenAiChatModel getOpenAiChatModel(String modelId) {
        return modelConfigMap.get(modelId);
    }

    @PostConstruct
    public void init() {
        log.info("Initializing OpenAI model configurations...");
        // 初始化模型配置
        List<ModelConfigDO> modelConfigDOS = modelConfigMapper.selectList(new LambdaQueryWrapper<ModelConfigDO>()
                .eq(ModelConfigDO::getChatClientType, ChatClientType.OpenAi)
        );
        modelConfigDOS.forEach(x -> {
            try {
                OpenAiApi build = OpenAiApi.builder()
                        .apiKey(x.getModelApikey())
                        .baseUrl(x.getModelEndpoint())
                        // 设置 completionsPath 为空则使用默认值 /v1/chat/completions
                        .completionsPath(x.getCompletionsPath())
                        .build();
                OpenAiChatOptions chatOptions = OpenAiChatOptions.builder().model(x.getModelName())
                        .streamUsage( true)
                        .build();
                OpenAiChatModel model = OpenAiChatModel.builder()
                        .openAiApi(build)
                        .defaultOptions(chatOptions)
                        .build();
                modelConfigMap.put(x.getId().toString(), model);
            } catch (Exception e) {
                log.error("Error in getModelById: ", e);
                throw new RuntimeException("Failed to get model by ID", e);
            }
        });
    }


}
